package org.xiaoyu.rpc.loadbalance;

import org.xiaoyu.rpc.Loadbalancer;
import org.xiaoyu.rpc.common.url.URL;
import org.xiaoyu.rpc.common.util.Bytes;
import org.xiaoyu.rpc.remoting.packet.Invocation;

import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.regex.Pattern;

/**
 * Created By Have
 * 2021/7/14 17:18
 */
public class ConsistentHashLoadbalancer implements Loadbalancer {
    private final ConcurrentMap<String, ConsistentHashSelector> selectors = new ConcurrentHashMap<>();
    private static final Pattern COMMA_SPLIT_PATTERN = Pattern.compile("\\s*[,]+\\s*");

    @Override
    public URL loadBalance(List<URL> urls, Invocation invocation) {

        String key = urls.get(0).getServiceName() + "." + invocation.getMethodName();
        int urlHashcode = urls.hashCode();
        ConsistentHashSelector selector = selectors.get(key);
        if (selector == null || urlHashcode != selector.hashCode()) {
            selectors.put(key, new ConsistentHashSelector(urls, urlHashcode, invocation.getMethodName()));
            selector = selectors.get(key);
        }
        return selector.select(invocation);
    }

    private static final class ConsistentHashSelector {
        private final TreeMap<Long, URL> virtualUrls;
        private final int identityHashCode;
        private final int replicaNumber;
        private final String methodName;
        private final int[] argumentIndex;
        public ConsistentHashSelector(List<URL> urls, int identityHashCode, String methodName) {
            this.virtualUrls = new TreeMap<>();
            this.identityHashCode = identityHashCode;
            this.replicaNumber = 160;
            this.methodName = methodName;
            String[] index = COMMA_SPLIT_PATTERN.split("0");
            argumentIndex = new int[index.length];
            for (int i = 0; i < index.length; i++) {
                argumentIndex[i] = Integer.parseInt(index[i]);
            }

            for (URL url : urls) {
                String host = url.getHost();
                int port = url.getPort();
                String address = host + ":" + port;
                for (int i = 0; i < replicaNumber / 4; i++) {
                    byte[] digest = Bytes.getMD5(address + i);
                    for (int h = 0; h < 4; h++) {
                        long m = hash(digest, h);
                        virtualUrls.put(m, url);
                    }
                }
            }
        }

        public URL select(Invocation invocation) {
            String key = toToken(new String[]{"0"});
            byte[] digest = Bytes.getMD5(key);
            return selectForKey(hash(digest, 0));
        }

        private URL selectForKey(long hash) {
            Map.Entry<Long, URL> entry = virtualUrls.ceilingEntry(hash);
            if (entry == null) {
                entry = virtualUrls.firstEntry();
            }
            return entry.getValue();
        }

        private String toToken(Object[] args) {
            StringBuilder sb = new StringBuilder();
            for (int index : argumentIndex) {
                if (index >= 0 && index < args.length) {
                    sb.append(args[index]);
                }
            }
            return sb.toString();
        }

        private long hash(byte[] digest, int number) {
            return (((long) (digest[3 + number * 4] & 0xFF) << 24)
                    | ((long) (digest[2 + number * 4] & 0xFF) << 16)
                    | ((long) (digest[1 + number * 4] & 0xFF) << 8)
                    | (digest[number * 4] & 0xFF))
                    & 0xFFFFFFFFL;
        }

        @Override
        public int hashCode() {
            return identityHashCode;
        }
    }
}
